from funasr import AutoModel
from funasr.utils.postprocess_utils import rich_transcription_postprocess
import time
from fastapi import APIRouter, UploadFile, File, Form, HTTPException
from pydantic import BaseModel
from typing import Optional
import os
import uuid
import jieba
import sys
import text_string_nlp_deal 
from config import *

# 创建缓存目录
cache_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), asr_cache_dir)
os.makedirs(cache_dir, exist_ok=True)
# 设置 jieba 缓存路径
jieba.dt.cache_file = os.path.join(cache_dir, "jieba.cache")
asr_router=APIRouter()
# 全局变量存储预加载的模型
global_model = None

def check_model_cache(model_name, cache_dir):
    """
    检查模型是否已经在默认缓存目录中（注意：模型缓存和字典缓存不是一个目录）
    
    Args:
        model_name: 模型名称
        cache_dir: 缓存目录
        
    Returns:
        bool: 模型是否已缓存
    """
    # 检查模型缓存目录
    model_cache_dir = os.path.join(cache_dir, "models")
    if not os.path.exists(model_cache_dir):
        return False
        
    # 处理模型名称，将路径分隔符替换为下划线
    cache_name = model_name.replace("/", "_")
    model_path = os.path.join(model_cache_dir, cache_name)
    
    # 检查模型目录是否存在
    if os.path.exists(model_path) and os.path.isdir(model_path):
        # 检查目录中是否有文件
        files = os.listdir(model_path)
        if len(files) > 0:
            return True
    
    return False

def load_model(model_dir=None, punc_model_dir=None, vad_model_dir=None, device="cuda:0"):
    """
    加载语音识别模型
    
    Args:
        model_dir: 模型目录或模型名称
        punc_model_dir: 标点符号模型目录或名称
        vad_model_dir: VAD模型目录或名称
        device: 使用的设备，例如"cuda:0"或"cpu"
        
    Returns:
        加载好的模型
    """
    global global_model
    
    try:
        # 使用配置文件中的路径（如果未提供）
        if model_dir is None:
            model_dir = asr_model_dir
        if punc_model_dir is None:
            punc_model_dir = asr_punc_model_dir
        if vad_model_dir is None:
            vad_model_dir = asr_vad_model_dir
        
        # 检查模型目录是否存在，如果不存在则使用备用模型
        if not os.path.exists(model_dir) and not model_dir.startswith("iic/"):
            # 检查备用模型是否已缓存
            if check_model_cache(asr_model_dir_backup, cache_dir):
                print(f"本地模型目录不存在: {model_dir}，使用已缓存的备用模型: {asr_model_dir_backup}")
            else:
                print(f"本地模型目录不存在: {model_dir}，尝试使用备用模型: {asr_model_dir_backup}")
            model_dir = asr_model_dir_backup
            
        # 检查标点模型目录是否存在，如果不存在则使用备用模型
        if not os.path.exists(punc_model_dir) and not punc_model_dir.startswith("iic/"):
            # 检查备用模型是否已缓存
            if check_model_cache(asr_punc_model_dir_backup, cache_dir):
                print(f"本地标点模型目录不存在: {punc_model_dir}，使用已缓存的备用模型: {asr_punc_model_dir_backup}")
            else:
                print(f"本地标点模型目录不存在: {punc_model_dir}，尝试使用备用模型: {asr_punc_model_dir_backup}")
            punc_model_dir = asr_punc_model_dir_backup
            
        # 检查VAD模型目录是否存在，如果不存在则使用备用模型
        if not os.path.exists(vad_model_dir) and not vad_model_dir.startswith("iic/"):
            # 检查备用模型是否已缓存
            if check_model_cache(asr_vad_model_dir_backup, cache_dir):
                print(f"本地VAD模型目录不存在: {vad_model_dir}，使用已缓存的备用模型: {asr_vad_model_dir_backup}")
            else:
                print(f"本地VAD模型目录不存在: {vad_model_dir}，尝试使用备用模型: {asr_vad_model_dir_backup}")
            vad_model_dir = asr_vad_model_dir_backup
        
        # 检测CUDA是否可用
        try:
            import torch
            cuda_available = torch.cuda.is_available()
            if cuda_available:
                device_to_use = device  # 使用传入的CUDA设备
                print(f"检测到CUDA可用，将使用 {device_to_use} 设备")
            else:
                device_to_use = "cpu"
                print("未检测到CUDA，将使用CPU进行推理")
        except ImportError:
            # 如果无法导入torch，默认使用CPU
            device_to_use = "cpu"
            print("无法导入PyTorch，将使用CPU进行推理")
        
        print(f"正在加载模型，使用以下配置:")
        print(f"- 模型: {model_dir}")
        print(f"- 标点模型: {punc_model_dir}")
        print(f"- VAD模型: {vad_model_dir}")
        print(f"- 设备: {device_to_use}")
        print(f"- 缓存目录: {cache_dir}")
        
        # 检查模型是否已缓存
        model_cached = check_model_cache(model_dir, cache_dir)
        punc_cached = check_model_cache(punc_model_dir, cache_dir)
        vad_cached = check_model_cache(vad_model_dir, cache_dir)
        
        if model_cached:
            print(f"模型已缓存: {model_dir}")
        if punc_cached:
            print(f"标点模型已缓存: {punc_model_dir}")
        if vad_cached:
            print(f"VAD模型已缓存: {vad_model_dir}")
        
        start_time = time.time()
        
        try:
            # 使用全局缓存目录，并设置 download 参数
            # 如果模型已缓存，则设置 download=False 避免重新下载
            model = AutoModel(
                model=model_dir,
                vad_model=vad_model_dir,
                vad_kwargs={"max_single_segment_time": 30000},
                device=device_to_use, 
                punc_model=punc_model_dir,
                disable_update=True,
                cache_dir=cache_dir,
                download=None  # None表示如果模型已存在则不下载，否则下载
            )
            
            end_time = time.time()
            print(f"模型加载完成，耗时 {end_time - start_time:.2f} 秒")
            
            # 更新全局模型
            global_model = model
            return model
            
        except Exception as e:
            print(f"加载模型时出错: {str(e)}")
            
            # 尝试使用所有备用模型
            print("尝试使用所有备用模型...")
            
            # 检查备用模型是否已缓存
            backup_model_cached = check_model_cache(asr_model_dir_backup, cache_dir)
            backup_punc_cached = check_model_cache(asr_punc_model_dir_backup, cache_dir)
            backup_vad_cached = check_model_cache(asr_vad_model_dir_backup, cache_dir)
            
            if backup_model_cached:
                print(f"备用模型已缓存: {asr_model_dir_backup}")
            if backup_punc_cached:
                print(f"备用标点模型已缓存: {asr_punc_model_dir_backup}")
            if backup_vad_cached:
                print(f"备用VAD模型已缓存: {asr_vad_model_dir_backup}")
            
            try:
                model = AutoModel(
                    model=asr_model_dir_backup,
                    vad_model=asr_vad_model_dir_backup,
                    vad_kwargs={"max_single_segment_time": 30000},
                    device=device_to_use, 
                    punc_model=asr_punc_model_dir_backup,
                    disable_update=True,
                    cache_dir=cache_dir,
                    download=None  # None表示如果模型已存在则不下载，否则下载
                )
                
                end_time = time.time()
                print(f"使用备用模型加载成功，耗时 {end_time - start_time:.2f} 秒")
                
                # 更新全局模型
                global_model = model
                return model
            except Exception as backup_error:
                print(f"使用备用模型加载失败: {str(backup_error)}")
                raise
            
    except Exception as e:
        print(f"模型加载失败: {str(e)}")
        raise

def transcribe_audio(audio_path, model_instance=None, model_dir=None, punc_model_dir=None, vad_model_dir=None, device="cuda:0", language="zh", use_itn=True, batch_size_s=10, merge_vad=True, merge_length_s=0.5,hotwords=asr_hotwords):
    """
    转写音频文件
    
    Args:
        audio_path: 音频文件路径
        model_instance: 已加载的模型实例，如果为None则会加载模型
        model_dir: 模型目录或模型名称
        punc_model_dir: 标点符号模型目录或名称
        vad_model_dir: VAD模型目录或名称
        device: 使用的设备，例如"cuda:0"或"cpu"
        language: 语言，例如"zh"或"en"
        use_itn: 是否使用inverse text normalization
        batch_size_s: 批处理大小（秒）
        merge_vad: 是否合并VAD结果
        merge_length_s: 合并长度（秒）
        
    Returns:
        转写结果和推理时间
    """
    global global_model
    
    # 检查音频文件是否存在
    if not os.path.exists(audio_path):
        raise FileNotFoundError(f"音频文件不存在: {audio_path}")
    
    # 检查音频文件是否为空
    if os.path.getsize(audio_path) == 0:
        raise ValueError(f"音频文件为空: {audio_path}")
    
    # 如果没有提供模型实例，尝试使用全局模型或加载模型
    if model_instance is None:
        if global_model is not None:
            model_instance = global_model
            print("使用已加载的全局模型")
        else:
            print("未提供模型实例且全局模型不可用，尝试加载模型...")
            try:
                # 检查模型是否已缓存
                if model_dir and check_model_cache(model_dir, cache_dir):
                    print(f"模型已缓存: {model_dir}")
                if punc_model_dir and check_model_cache(punc_model_dir, cache_dir):
                    print(f"标点模型已缓存: {punc_model_dir}")
                if vad_model_dir and check_model_cache(vad_model_dir, cache_dir):
                    print(f"VAD模型已缓存: {vad_model_dir}")
                
                model_instance = load_model(model_dir, punc_model_dir, vad_model_dir, device)
            except Exception as e:
                print(f"加载模型失败，尝试使用备用模型: {str(e)}")
                try:
                    # 检查备用模型是否已缓存
                    if check_model_cache(asr_model_dir_backup, cache_dir):
                        print(f"备用模型已缓存: {asr_model_dir_backup}")
                    if check_model_cache(asr_punc_model_dir_backup, cache_dir):
                        print(f"备用标点模型已缓存: {asr_punc_model_dir_backup}")
                    if check_model_cache(asr_vad_model_dir_backup, cache_dir):
                        print(f"备用VAD模型已缓存: {asr_vad_model_dir_backup}")
                    
                    model_instance = load_model(
                        asr_model_dir_backup, 
                        asr_punc_model_dir_backup, 
                        asr_vad_model_dir_backup, 
                        device
                    )
                except Exception as backup_error:
                    raise RuntimeError(f"加载备用模型也失败: {str(backup_error)}")
    
    # 执行转写
    try:
        start_time = time.time()
        
        # 执行转写，不使用缓存参数以避免 prev_samples 错误
        result = model_instance.generate(
            input=audio_path,
            language=language,
            use_itn=use_itn,
            batch_size_s=batch_size_s,
            vad_merge=merge_vad,
            vad_merge_length_s=merge_length_s,
            hotwords=hotwords,
            hotword_weight=10.0
        )
        
        end_time = time.time()
        inference_time = end_time - start_time
        
        print(f"原始转写结果类型: {type(result)}")
        print(f"原始转写结果内容: {result}")
        
        # 获取转写文本 - 根据 FunASR 返回的可能格式处理
        if isinstance(result, list):
            if len(result) > 0:
                if isinstance(result[0], dict) and "text" in result[0]:
                    # 使用 rich_transcription_postprocess 处理文本
                    transcribed_text = rich_transcription_postprocess(result[0]["text"])
                    # 去除<>以及中括号内的内容
                    if text_string_nlp_deal:
                        transcribed_text = text_string_nlp_deal.text_rm_brack_chinese_marks_punctuation_replace(transcribed_text)
                else:
                    # 尝试将列表中的第一项作为文本
                    transcribed_text = str(result[0])
            else:
                transcribed_text = "转写结果为空"
        else:
            # 如果不是列表，直接使用返回结果
            transcribed_text = str(result)
        
        print(f"处理后的转写结果: {transcribed_text}")
        
        return {
            "text": transcribed_text,
            "inference_time": inference_time
        }
    
    except Exception as e:
        print(f"转写音频时出错: {str(e)}")
        raise

class TranscriptionResponse(BaseModel):
    """转录响应模型"""
    text: str  # 明确声明为字符串类型
    duration: float  # 推理时间，浮点数类型

@asr_router.post("/transcribe", response_model=TranscriptionResponse)
async def transcribe_endpoint(
    audio_file: UploadFile = File(...),
    language: str = Form("auto"),
    use_itn: bool = Form(True),
    batch_size_s: int = Form(60),
    merge_vad: bool = Form(True),
    merge_length_s: int = Form(15)
):
    """
    接收音频文件并返回转录结果的API端点
    """
    global global_model
    
    # 检查文件格式
    valid_extensions = ['.wav', '.mp3', '.flac', '.ogg', '.m4a', '.aac', '.wma']
    file_extension = os.path.splitext(audio_file.filename)[1].lower() if audio_file.filename else ""
    if not file_extension or file_extension not in valid_extensions:
        raise HTTPException(status_code=400, detail=f"不支持的音频格式: {file_extension}。支持的格式: {', '.join(valid_extensions)}")
    
    # 确保模型已加载
    if global_model is None:
        try:
            global_model = load_model()
        except Exception as e:
            raise HTTPException(status_code=500, detail=f"加载模型失败: {str(e)}")
    
    # 创建临时文件保存上传的音频
    temp_dir = "temp"
    os.makedirs(temp_dir, exist_ok=True)
    
    temp_path = os.path.join(temp_dir, f"{uuid.uuid4()}{file_extension}")
    
    try:
        # 保存上传的音频文件
        content = await audio_file.read()
        
        # 检查文件是否为空
        if len(content) == 0:
            raise HTTPException(status_code=400, detail="上传的音频文件为空")
            
        with open(temp_path, "wb") as f:
            f.write(content)
        
        # 检查文件是否正确保存
        if not os.path.exists(temp_path) or os.path.getsize(temp_path) == 0:
            raise HTTPException(status_code=500, detail="无法保存上传的音频文件")
            
        # 转录音频
        print(f"开始转录音频: {temp_path}")
        start_time = time.time()
        try:
            transcription_result = transcribe_audio(
                audio_path=temp_path,
                model_instance=global_model,  # 直接传递模型实例
                language=language,
                use_itn=use_itn,
                batch_size_s=batch_size_s,
                merge_vad=merge_vad,
                merge_length_s=merge_length_s
            )
            print(f"转录结果: {transcription_result}")
        except Exception as e:
            # 提供更详细的错误信息
            error_message = f"音频转录失败: {str(e)}"
            print(error_message)
            raise HTTPException(status_code=500, detail=error_message)
        
        # 从返回结果中获取文本和推理时间
        if isinstance(transcription_result, dict):
            if "text" in transcription_result and "inference_time" in transcription_result:
                text = transcription_result["text"]
                inference_time = transcription_result["inference_time"]
            else:
                error_message = f"返回的转录结果格式不正确: {transcription_result}"
                print(error_message)
                raise HTTPException(status_code=500, detail=error_message)
        else:
            error_message = f"返回的转录结果类型不正确: {type(transcription_result)}"
            print(error_message)
            raise HTTPException(status_code=500, detail=error_message)
        
        # 确保文本是字符串类型
        if not isinstance(text, str):
            text = str(text)
        
        print(f"最终返回文本: {text} (类型: {type(text)})")
        print(f"推理时间: {inference_time} 秒")
        
        # 返回转录结果
        return TranscriptionResponse(text=text, duration=inference_time)
    
    except HTTPException:
        # 重新抛出 HTTP 异常
        raise
    except Exception as e:
        # 捕获并记录所有其他异常
        error_message = f"转录失败: {str(e)}"
        print(error_message)
        raise HTTPException(status_code=500, detail=error_message)
    
    finally:
        # 删除临时文件
        try:
            if os.path.exists(temp_path):
                os.remove(temp_path)
        except Exception as e:
            print(f"删除临时文件失败: {str(e)}")

# 程序启动时预加载模型
if __name__ == "__main__":
    # 预加载模型，自动使用配置文件中的路径
    global_model = load_model()
    
    # 示例：转录英文音频
    example_audio = f"{global_model.model_path}/example/en.mp3"
    transcription = transcribe_audio(audio_path=example_audio, model_instance=global_model)
    print("\n转录结果:")
    print(transcription)
    
    # 如果需要转录其他音频文件，可以直接调用，不需要重新加载模型
    # 例如：
    # another_audio = "path/to/another/audio.wav"
    # another_transcription = transcribe_audio(audio_path=another_audio, model_instance=global_model)
    # print(another_transcription)